import socket
import threading
import sqlite3
import time
import logging
import signal

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='[%(asctime)s] %(message)s',
    handlers=[
        logging.FileHandler("server.log"),
        logging.StreamHandler()
    ]
)

HEARTBEAT_TIMEOUT = 30  # 心跳超时时间


# 初始化数据库
def init_db():
    conn = sqlite3.connect('license_server.db')
    cursor = conn.cursor()
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS licenses (
        serial_number TEXT PRIMARY KEY,
        username TEXT,
        password TEXT,
        license_type TEXT,
        max_users INTEGER,
        current_users INTEGER
    )
    ''')
    conn.commit()
    conn.close()


# 状态保存
def save_state():
    conn = sqlite3.connect('license_server.db')
    cursor = conn.cursor()
    cursor.execute('SELECT serial_number, current_users FROM licenses')
    licenses = cursor.fetchall()
    conn.close()
    logging.info(f"State saved: {licenses}")


# 状态加载
def load_state():
    conn = sqlite3.connect('license_server.db')
    cursor = conn.cursor()
    cursor.execute('SELECT serial_number, current_users FROM licenses')
    licenses = cursor.fetchall()
    conn.close()
    logging.info(f"State loaded: {licenses}")


# 生成许可证
def generate_license(username, password, license_type, max_users):
    import random
    serial_number = ''.join([str(random.randint(0, 9)) for _ in range(10)])
    conn = sqlite3.connect('license_server.db')
    cursor = conn.cursor()
    cursor.execute(
        'INSERT INTO licenses (serial_number, username, password, license_type, max_users, current_users) VALUES (?, ?, ?, ?, ?, ?)',
        (serial_number, username, password, license_type, max_users, 0))
    conn.commit()
    conn.close()
    return serial_number


# 捕获 SIGINT 信号
def signal_handler(sig, frame):
    logging.info("Server is shutting down. Saving state...")
    save_state()
    exit(0)


# 处理客户端
def handle_client(client_socket, client_address):
    logging.info(f"Connection from {client_address}")
    conn = sqlite3.connect('license_server.db')
    cursor = conn.cursor()
    try:
        last_heartbeat = time.time()
        serial_number = None

        while True:
            client_socket.settimeout(HEARTBEAT_TIMEOUT)
            try:
                data = client_socket.recv(1024).decode()
                if not data:
                    raise ConnectionResetError("Client closed connection")

                command, *args = data.split(' ')
                if command == 'VALIDATE':
                    serial_number = args[0]
                    cursor.execute('SELECT max_users, current_users FROM licenses WHERE serial_number = ?',
                                   (serial_number,))
                    result = cursor.fetchone()
                    if result:
                        max_users, current_users = result
                        if current_users < max_users:
                            cursor.execute(
                                'UPDATE licenses SET current_users = current_users + 1 WHERE serial_number = ?',
                                (serial_number,))
                            conn.commit()
                            logging.info(f"License validated for {serial_number}. Current users: {current_users + 1}")
                            client_socket.send(b'AUTHORIZED')
                        else:
                            client_socket.send(b'DENIED')
                    else:
                        client_socket.send(b'INVALID')
                elif command == 'HEARTBEAT':
                    serial_number = args[0]
                    last_heartbeat = time.time()
                    logging.info(f"Heartbeat received from {client_address} for license {serial_number}")
                    client_socket.send(b'ALIVE')
                elif command == 'EXIT':
                    serial_number = args[0]
                    cursor.execute('UPDATE licenses SET current_users = current_users - 1 WHERE serial_number = ?',
                                   (serial_number,))
                    conn.commit()
                    logging.info(f"Client {client_address} exited normally. Reduced license usage.")
                    client_socket.send(b'DISCONNECTED')
                    break

                if time.time() - last_heartbeat > HEARTBEAT_TIMEOUT:
                    raise TimeoutError("Heartbeat timeout")

            except (socket.timeout, TimeoutError):
                logging.warning(
                    f"Client {client_address} lost connection due to heartbeat timeout. Reducing license usage.")
                if serial_number:
                    cursor.execute('UPDATE licenses SET current_users = current_users - 1 WHERE serial_number = ?',
                                   (serial_number,))
                    conn.commit()
                break
            except (ConnectionResetError, socket.error):
                logging.error(f"Client {client_address} disconnected unexpectedly. Reducing license usage.")
                if serial_number:
                    cursor.execute('UPDATE licenses SET current_users = current_users - 1 WHERE serial_number = ?',
                                   (serial_number,))
                    conn.commit()
                break

    except Exception as e:
        logging.error(f"Error handling client {client_address}: {e}")
    finally:
        client_socket.close()


# 启动服务器
def start_server(host='0.0.0.0', port=8080):
    init_db()
    load_state()
    signal.signal(signal.SIGINT, signal_handler)
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.bind((host, port))
    server_socket.listen(5)
    logging.info("License server started...")
    logging.info("Use 'python3 server.py admin' to create new licenses.")

    while True:
        client_socket, client_address = server_socket.accept()
        threading.Thread(target=handle_client, args=(client_socket, client_address)).start()


# 管理员界面
def admin_interface():
    print("Admin Interface:")
    username = input("Enter username: ")
    password = input("Enter password: ")
    license_type = input("Enter license type (basic/pro): ")
    max_users = int(input("Enter maximum number of users: "))
    serial_number = generate_license(username, password, license_type, max_users)
    print(f"Generated License Key: {serial_number}")


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "admin":
        admin_interface()
    else:
        start_server()
